import argparse
import os  
import acl   # type: ignore
import numpy as np  
from PIL import Image
import cv2  
import time

ACL_MEM_MALLOC_HUGE_FIRST = 0  
ACL_MEMCPY_HOST_TO_DEVICE = 1  
ACL_MEMCPY_DEVICE_TO_HOST = 2  
  
def images_preprocess(image_path,input_size,bs):
    img = cv2.imread(image_path)
    img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
    img = cv2.resize(img, input_size)
    img = img.transpose(2,0,1)
    img = img.astype(np.float32)
    return img
  
class Net:  
    def __init__(self, model_path, device:0):   
        self.device_id = device   
        ret = acl.init()  
        ret = acl.rt.set_device(self.device_id)  
        self.model_id, ret = acl.mdl.load_from_file(model_path)  
        self.model_desc = acl.mdl.create_desc()  
        ret = acl.mdl.get_desc(self.model_desc, self.model_id) 
        self.input_dataset, self.input_data = self.creaetDataset('input')
        self.output_dataset, self.output_data = self.creaetDataset('output')

    def creaetDataset(self, io_type): 
        if io_type == "input": 
            io_num = acl.mdl.get_num_inputs(self.model_desc) 
            acl_mdl_get_size_by_index = acl.mdl.get_input_size_by_index
        else:
            io_num = acl.mdl.get_num_outputs(self.model_desc) 
            acl_mdl_get_size_by_index = acl.mdl.get_output_size_by_index  
        dataset = acl.mdl.create_dataset()
        datas = []
        for i in range(io_num):  
            buffer_size = acl_mdl_get_size_by_index(self.model_desc, i)  
            buffer, ret = acl.rt.malloc(buffer_size, ACL_MEM_MALLOC_HUGE_FIRST)  
            data_buffer = acl.create_data_buffer(buffer, buffer_size)  
            _, ret = acl.mdl.add_dataset_buffer(dataset, data_buffer)  
            datas.append({"buffer": buffer, "data": data_buffer, "size": buffer_size})  

        return dataset, datas 
    
    def forward(self, inputs): 
        input_num = len(inputs)
        for i in range(input_num):  
            bytes_data = inputs[i].tobytes()  
            bytes_ptr = acl.util.bytes_to_ptr(bytes_data)    
            ret = acl.rt.memcpy(self.input_data[i]["buffer"],  
                                self.input_data[i]["size"],  
                                bytes_ptr,   
                                len(bytes_data),   
                                ACL_MEMCPY_HOST_TO_DEVICE)  
        ret = acl.mdl.execute(self.model_id, self.input_dataset, self.output_dataset)    
        inference_result = []  
        for i, item in enumerate(self.output_data):  
            buffer_host, ret = acl.rt.malloc_host(self.output_data[i]["size"])   
            ret = acl.rt.memcpy(buffer_host,   
                                self.output_data[i]["size"], 
                                self.output_data[i]["buffer"],  
                                self.output_data[i]["size"],  
                                ACL_MEMCPY_DEVICE_TO_HOST)  
            bytes_out = acl.util.ptr_to_bytes(buffer_host, self.output_data[i]["size"]) 
            data = np.frombuffer(bytes_out, dtype=np.float32)
            inference_result.append(data)

        return inference_result
    
    def __del__(self):  
        # for dataset in [self.input_data, self.output_data]:  
        #     while dataset:  
        #         item = dataset.pop()  
        #         ret = acl.destroy_data_buffer(item["data"])  
        #         ret = acl.rt.free(item["buffer"]) 
        ret = acl.mdl.destroy_dataset(self.input_dataset)  
        ret = acl.mdl.destroy_dataset(self.output_dataset)
        ret = acl.mdl.destroy_desc(self.model_desc)  
        ret = acl.mdl.unload(self.model_id)  
        ret = acl.rt.reset_device(self.device_id)  
        ret = acl.finalize()

def iou(box1,box2):
    x1,y1,w1,h1 = box1
    x2,y2,w2,h2 = box2

    _x1 = max(x1,x2)
    _x2 = min(x1+w1,x2+w2)
    _y1 = max(y1,y2)
    _y2 = min(y1+h1,y2+h2)

    _area = max(0, _x2 - _x1) * max(0, _y2 - y1)
    area1 = w1*h1
    area2 = w2*h2
    union_area = area1 + area2 -_area

    return _area / union_area

def nms(boxs,scores,iou_threshold=0.5):

    selected_indexs = []
    while len(boxs) > 0:
        max_indx = np.argmax(scores)
        selected_indexs.append(max_indx)

        boxs = np.delete(boxs,max_indx,axis=0)
        scores = np.delete(scores,max_indx,axis=0)

        if len(boxs) > 0:
            ious = np.array([iou(boxs[max_indx],box) for box in boxs])
            boxs = boxs[ious <= iou_threshold]
            scores = scores[ious <= iou_threshold]

    return selected_indexs


def post_process(output,num_classes,conf=0.5,nmsflag=False,iouf=0.5):
    boxes = output[:4,:].T # [8400,4]
    scores = output[4:4+num_classes,:].T # [8400,80]
    breakpoint()
    class_ids = np.argmax(scores,axis=1) # 8400,
    confidences = np.max(scores,axis=1) # 8400,

    val_index = np.where(confidences > conf)[0] if not nmsflag else nms(boxes,confidences,iouf)
    boxes = boxes[val_index]
    class_ids = class_ids[val_index]
    confidences = confidences[val_index]
    results = []
    for i in range(len(boxes)):
        x,y,w,h = boxes[i]
        class_id =class_ids[i]
        confidence = confidences[i]
        results.append((x,y,w,h,class_id,confidence))

    return results


if __name__ == "__main__":  
    parser = argparse.ArgumentParser(description="yolov11 dataset convert")
    parser.add_argument("--om_path", type=str, help="weight")
    parser.add_argument("--data_path", type=str, help="dataset")

    args = parser.parse_args()
    model = Net(args.om_path,0) 
    image_list = [os.path.join(args.data_path,f) for f in os.listdir(args.data_path) if f.endswith('.jpg')]
    
    batch_size=16
    image_shape=(640,640)

    input_datas = []
    for i in range(0,len(image_list),batch_size):  
        batch_img = [images_preprocess(f,image_shape,batch_size) for f in image_list[i:i+batch_size]]
        input_data = np.stack(batch_img,axis=0)
        input_datas.append(input_data)


    start_time = time.time() * 1000
    for in_data in input_datas:  
        result = model.forward([in_data])
        # breakpoint()
        output = result[0].reshape(16,84,8400)
        a = post_process(output[0,:,:],80,0.5,False,0.5)

    end_time = time.time() * 1000

    diff_t = round((end_time -start_time),3)
    img_pre_time = diff_t / len(image_list) 
    print(f'images:{len(image_list)}, time:{diff_t}, img_pre_time:{img_pre_time}')

